from vtk import vtkSTLReader, vtkTextActor, vtkPoints, vtkLandmarkTransform, vtkRenderer, vtkPolyDataMapper, vtkSphereSource, vtkConeSource, vtkPolyDataReader, vtkPolyDataReader, vtkActor, vtkTriangleFilter, vtkPolyDataNormals, vtkPolyData
from RecordingState import RecordingState
from PlaybackState import PlaybackState
from Sensor import Sensor
from RecordingsProcessor import RecordingsProcessor
from numpy import *
from time import strftime
import wx
import os, sys

##  VisualizationState contains the functionality behind the VisualizationPanel.
#   This state ensures with for every frame the bones of the skeleton model are set between the calculated bony landmarks.
#   It also takes care of the playback of these frames.
class VisualizationState:

    ##  The constructor of VisualizationState.
    #   @param[in]  main    Main wx.App derived instance.
    def __init__(self, main):
        ##  Main wx.App derived instance.
        self.main = main
        ##  The framerate of the visualization.
        self.frameRate = self.main.config.ReadFramerate()        
        ##  Used to test between the Playback and Recording subpanel of VisualizationPanel.
        self.panel = None
        ##  Instantiation of the PlaybackState.
        self.playbackState = PlaybackState(self,main)
        ##  Instantiation of the RecordingState.
        self.recordingState = RecordingState(self, main)
        ##  Instantiation of the RecordingProcessor.
        self.recordingsProcessor = RecordingsProcessor(self, main)
        self.__meshPos = self.__GetMeshPositions()

    def get_actor_matrices(self):
        actorList = self.playbackState.actorList
        scap = actorList[81].GetUserTransform().GetMatrix() # scapula
        hum = actorList[82].GetUserTransform().GetMatrix() # humerus
        return [scap, hum]
        
        
        
    ##  Sets the state to the current visualization subpanel.
    #   Updates the bony landmarks to their current position, using the sensor positions.
    #   Also takes care of the rendering of the current position.
    def update_actor_positions(self):
        if self.panel.visNotebook.GetCurrentPage().name=="Playback" and self.playbackState.selectedRecording != None:
            state = self.playbackState
            recording = state.selectedRecording
            sensorsUpdate = recording.GetSensorsCurrentFrame()
            bonyLandmarksUpdate = recording.bonyLandmarksRecorded
            timeStampUpdate = recording.GetCurrentTimeStamp()
        else:
            state = self.recordingState
            sensorsUpdate = self.main.sensors
            bonyLandmarksUpdate = self.main.bonyLandmarks
            timeStampUpdate = strftime("%Y-%m-%d %H:%M:%S")
        # Update the time stamp display.
        #state.timeStampActor.SetInput(timeStampUpdate)

        ## List containing the current bony landmark positions in the visualization.
        self.__blPos = []

        for i in range(0, len(bonyLandmarksUpdate)):
            # This bony landmark has a sensor that determines its position.
            if bonyLandmarksUpdate[i].ID != -1:
                state.blActorList[i].AddPosition(
                                                 -state.blActorList[i].GetCenter()[0],
                                                 -state.blActorList[i].GetCenter()[1],
                                                 -state.blActorList[i].GetCenter()[2]
                                                 )

                relativeSensor = None
                for sensor in sensorsUpdate:
                    if sensor.ID == bonyLandmarksUpdate[i].sensorID:
                        relativeSensor = sensor

                blPos = self.main.driver.RelativeToGlobal(bonyLandmarksUpdate[i].relativePosition,
                                                          relativeSensor.position,
                                                          relativeSensor.rotation)

                state.blActorList[i].AddPosition(blPos[0], blPos[1], blPos[2])
                self.__blPos.append(blPos)

        # Update sensors.
        for i in range(0, len(sensorsUpdate)):
            state.sensorActorList[i].AddPosition(
                                                 -state.sensorActorList[i].GetCenter()[0],
                                                 -state.sensorActorList[i].GetCenter()[1],
                                                 -state.sensorActorList[i].GetCenter()[2]
                                                 )
            state.sensorActorList[i].AddPosition(
                                                 sensorsUpdate[i].position[0],
                                                 sensorsUpdate[i].position[1],
                                                 sensorsUpdate[i].position[2]
                                                 )
        self.SetBones()

        if state.FIRST_TIME:
            state.ren.ResetCamera()
            self.InitCamera(state)
            state.FIRST_TIME = False

        if self.main.motion_analysis_started:
            if self.__blPos[12].tolist() != self.__blPos[13].tolist()\
                and self.__blPos[22].tolist() != self.__blPos[23].tolist():
                blPosList = []
                for i in range(0, len(self.__blPos)):
                    blPosList.append( [self.__blPos[i].tolist()[0][0],\
                                               self.__blPos[i].tolist()[1][0],\
                                               self.__blPos[i].tolist()[2][0]] )
                                                                           
                self.main.angleCalculations.SetupCoordinateSystems(blPosList)
                self.main.angleCalculations.Angles()
            
        state.panel.widget.Render()

    ##  Maps every bone of the skeletal model to the actual bony landmark coordinates.
    def SetBones(self):
        # Read the 3d model bones and bony landmark connection from the config file.
        modelConnectionList =  self.main.config.Read3DModelConnection()

        # Call the SetBones method for each connection.
        for i in range(0, len(modelConnectionList)):
            self.SetBone(modelConnectionList[i][0], modelConnectionList[i][1])

    ##  Maps coordinates of the bony landmarks in the model to the actual coordinates of the bony landmarks.
    #   @param[in]  boneIndices The bone to be mapped to the bony landmark positions.
    #   @param[in]  blIndices   The bony landmark.
    def SetBone(self, boneIndices, blIndices):
        if self.panel.visNotebook.GetCurrentPage().name=="Playback" and self.playbackState.selectedRecording != None:
            actorList = self.playbackState.actorList
            state = self.playbackState
            bonyLandmarksUpdate = self.playbackState.selectedRecording.bonyLandmarksRecorded
        else:
            actorList = self.recordingState.actorList
            state = self.recordingState
            bonyLandmarksUpdate = self.main.bonyLandmarks

        for j in range(0, len(boneIndices)):
            if self.main.gui.Is3DModelChecked():
                actorList[boneIndices[j]].VisibilityOn()

        mPoints = vtkPoints()
        rPoints = vtkPoints()

        # Because the bony landmark ID won't be the same as the index anymore, the blIDToIndex is used to map the bony landmark ID to its index.
        # This was done to improve on the generic aspects of the program.
        blIDToIndex = self.main.GetBonyLandmarkDict(True)

        for i in range(0, len(blIndices)):
            ID = blIndices[i]
            if self.main.gui.IsBonyLandmarksChecked():
                state.blActorList[blIDToIndex[ID]].VisibilityOn()
            # If there are no measurements available for the bony landmark, turn visibility off.
            if bonyLandmarksUpdate[blIDToIndex[ID]].relativePosition[0]==0.0 and \
               bonyLandmarksUpdate[blIDToIndex[ID]].relativePosition[1]==0.0 and \
               bonyLandmarksUpdate[blIDToIndex[ID]].relativePosition[2]==0.0:
                   for j in range(0, len(boneIndices)):
                       actorList[boneIndices[j]].VisibilityOff()
                   state.blActorList[blIDToIndex[ID]].VisibilityOff()
            # Don't use the bl ID to Index mapping for the meshPos. The Bony landmark ID corresponds with its position index in the blpositions.ini file.
            mPoints.InsertPoint(i, self.__meshPos[ID])
            rPoints.InsertPoint(i, self.__blPos[blIDToIndex[ID]])

        transformer = vtkLandmarkTransform()
        transformer.SetSourceLandmarks(mPoints)
        transformer.SetTargetLandmarks(rPoints)

        for i in range(0, len(boneIndices)):
            actorList[boneIndices[i]].SetUserTransform(transformer)

    ##  Toggles the visualization of the sensor on the model on and off.
    def ToggleSensors(self):
        visible = self.main.gui.IsSensorsChecked()

        for actor in self.playbackState.sensorActorList:
            actor.SetVisibility(visible)

        for actor in self.recordingState.sensorActorList:
            actor.SetVisibility(visible)

        try:
            self.playbackState.panel.widget.Render()
            self.recordingState.panel.widget.Render()
        except:
            pass

    ##  Toggles the visualization of the bony landmarks on the model on and off.
    def ToggleBonyLandmarks(self):
        visable = self.main.gui.IsBonyLandmarksChecked()

        for actor in self.playbackState.blActorList:
            actor.SetVisibility(visable)

        for actor in self.recordingState.blActorList:
            actor.SetVisibility(visable)

        try:
            self.playbackState.panel.widget.Render()
            self.recordingState.panel.widget.Render()
        except:
            pass

    ##  Toggles the visualization of the 3D the model on and off.
    def Toggle3DModel(self):
        visable = self.main.gui.Is3DModelChecked()

        for actor in self.playbackState.actorList:
            actor.SetVisibility(visable)

        for actor in self.recordingState.actorList:
            actor.SetVisibility(visable)

        try:
            self.playbackState.panel.widget.Render()
            self.recordingState.panel.widget.Render()
        except:
            pass

    ##  Method used for stopping animation. A check is done whether visualization is playing. 
    #   If so, it is stopped using onTogglePlayPause.
    def Stop(self):
        if self.playbackState.PLAYING:
            self.playbackState.OnTogglePlayPause()

        if self.recordingState.PLAYING:
            self.recordingState.OnTogglePlayPause()

    ##  Event handler, starts or stops the playback of the visualization in the opened VTK window.
    def OnTogglePlayPause(self):
        if self.panel.visNotebook.GetCurrentPage().name == "Recording":
            self.recordingState.OnTogglePlayPause()
        else:
            self.playbackState.OnTogglePlayPause()

    ##  Reads the mesh positions out of a file.
    #   @return[matrix] The mesh positions.
    def __GetMeshPositions(self):
        try:
            file = open(self.main.prefReader.blPositionsFile, 'r')
        except:
            msg = "Could not open %s." % self.main.prefReader.blPositionsFile
            dlg = wx.MessageDialog(self.main.gui, msg, 'Warning', wx.OK | wx.NO_DEFAULT | wx.ICON_EXCLAMATION)
            dlg.ShowModal()
        lines = file.read().split('\n')
        meshPos = []
        for line in lines:
            line = line.split()
            meshPos.append([float(x) for x in line])
        return meshPos

    ##  Initiates the VTK screen and the needed models and camera.
    #   @param[in]  state   An instantiation of VisualizationState, PlaybackState or RecordingState.
    def InitiateVTK(self, state):

        state.ren = vtkRenderer()
        state.ren.SetBackground(.92, .93, .95)

        state.ren.ResetCamera()
        state.ren.GetActiveCamera().Azimuth(90)
        state.ren.GetActiveCamera().Roll(90)

        ## Bony Landmark Actor List used for rendering bony landmarks spheres.
        state.blActorList = []
        ## Sensor Actor List used for rendering sensors spheres.
        state.sensorActorList = []
        ## Actor List used for rendering the 3d skeleton model.
        state.actorList = []

        ## Set the timeStamp Actor for the state, and the text format properties.
        state.timeStampActor = vtkTextActor()
        state.timeStampActor.SetTextScaleModeToViewport()
        state.timeStampActor.SetDisplayPosition(20,20)
        state.timeTextProperty = state.timeStampActor.GetTextProperty()
        state.timeTextProperty.SetColor(255,255,255)
        state.timeTextProperty.SetFontSize(7)
        state.timeStampActor.SetTextProperty(state.timeTextProperty)
        state.ren.AddActor(state.timeStampActor)

    ##  Sets the actors for the according state and its render window.
    #   Always called after VTK is initiated.
    #   @param[in]    state              The state using this method.
    #   @param[in]    numBLActors        The number of bony landmarks to be displayed by the renderer.
    #   @param[in]    numSensorActors    The number of Sensors to be displayed by the renderer.
    #   @param[in]    numGHJoints        The number of GH Joints the renderer works with.
    def SetRenderActors(self, state, numBLActors, numSensorActors, numGHJoints):

        ## Remove Actors if they have been added to the state renderer, re-add the timeStampActor.
        allActors = concatenate((state.blActorList, state.sensorActorList, state.actorList),1)
        for actor in allActors: state.ren.RemoveActor(actor)
        state.ren.AddActor(state.timeStampActor)

        blMapperList = []
        blSphereList = []

        sensorMapperList = []
        sensorSphereList = []
        for i in range(0,numBLActors):
            state.blActorList.append(vtkActor())
            blMapperList.append(vtkPolyDataMapper())
            blSphereList.append(vtkSphereSource())
            blMapperList[i].SetInput(blSphereList[i].GetOutput())
            state.blActorList[i].SetMapper(blMapperList[i])
            # Colour BLs Yellow
            state.blActorList[i].GetProperty().SetColor(1,1,0)
            state.ren.AddActor(state.blActorList[i])

        # Colour GH-Joints Green.
        # Assuming the last bony landmarks are the two GH-Joints that have been appended to the end of the bony landmark list.
        for i in range(1,numGHJoints+1):
            state.blActorList[len(state.blActorList)-i].GetProperty().SetColor(0,1,0)
            #state.blActorList[len(state.blActorList)-2].GetProperty().SetColor(0,1,0)
            #state.blActorList[len(state.blActorList)-1].GetProperty().SetColor(0,1,0)

        for i in range(0,numSensorActors):
            state.sensorActorList.append(vtkActor())
            sensorMapperList.append(vtkPolyDataMapper())
            sensorSphereList.append(vtkSphereSource())
            sensorMapperList[i].SetInput(sensorSphereList[i].GetOutput())
            state.sensorActorList[i].SetMapper(sensorMapperList[i])
            # Colour Sensors Red.
            state.sensorActorList[i].GetProperty().SetColor(1,0,0)
            if i != 0:
                state.ren.AddActor(state.sensorActorList[i])

        mapperList = []
        readerList = []
        triangleList = []
        normalList = []

        for i in range(0, self.main.config.Read3DModelParts()):

                state.actorList.append(vtkActor())
                mapperList.append(vtkPolyDataMapper())
                
                filename = self.main.config.Read3DFileName()
                for j in self.main.replacement_tuples:
                    if i == j[0]:                        
                        filename = j[1]
                        print 'Replacing: ', filename

                if os.path.splitext(filename)[-1]=='.vtk':
                    readerList.append(vtkPolyDataReader())
                    readerList[i].SetFileName(filename % (i+1))
                    #if i == 
                    readerList[i].OpenVTKFile()
                else:
                    readerList.append(vtkSTLReader())
                    readerList[i].SetFileName(filename)
                    #if i == 
                    readerList[i].Update()

                if self.main.Smooth3DModel:
                    readerList[i].Update()
                    p1 = readerList[i].GetOutput()
                    p2 = vtkPolyData()
                    p2.DeepCopy(p1)

                    triangleList.append(vtkTriangleFilter())
                    triangleList[i].SetInputConnection(readerList[i].GetOutputPort())
                    normalList.append(vtkPolyDataNormals())
                    normalList[i].SetInput(triangleList[i].GetOutput())
                    normalList[i].SplittingOff()
                    normalList[i].ConsistencyOff()
                    mapperList[i].SetInput(normalList[i].GetOutput())
                else:
                    mapperList[i].SetInput(readerList[i].GetOutput())

                mapperList[i].Update()
                state.actorList[i].SetMapper(mapperList[i])
                state.ren.AddActor(state.actorList[i])
                state.actorList[i].GetProperty().SetColor(1, .97, .85)
                state.actorList[i].VisibilityOff()
                
                if os.path.splitext(filename)[-1]=='.vtk':
                    readerList[i].CloseVTKFile()

    ## Initialises the camera position, used when it is reset.
    #  @param[in]   state   The state in which the camera will be initialized.
    def InitCamera(self,state):
        state.initCamPosition = state.ren.GetActiveCamera().GetPosition()
        state.initCamFocalPoint = state.ren.GetActiveCamera().GetFocalPoint()
        state.initCamViewUp = state.ren.GetActiveCamera().GetViewUp()

    ##  Resets the camera for the state passed.
    #   @param[in] state    The state in which the camera will be reset.
    def ResetCamera(self,state):
        state.ren.GetActiveCamera().SetPosition(state.initCamPosition)
        state.ren.GetActiveCamera().SetFocalPoint(state.initCamFocalPoint)
        state.ren.GetActiveCamera().SetViewUp(state.initCamViewUp)
        state.ren.ResetCamera()
    
    def UpdateBL(self, index, bl):
        self.__meshPos[index] = bl

    def GetBL(self, index):
        return self.__meshPos[index]